package org.truenewx.tnxjee.repo.mybatis.mapper.query;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.function.Function;

import com.baomidou.mybatisplus.core.MybatisMapperAnnotationBuilder;
import com.baomidou.mybatisplus.core.MybatisMapperRegistry;
import com.baomidou.mybatisplus.core.override.MybatisMapperProxyFactory;
import org.apache.ibatis.binding.MapperRegistry;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.truenewx.tnxjee.core.util.BeanUtil;
import org.truenewx.tnxjee.model.query.QueryResult;
import org.truenewx.tnxjee.model.query.Querying;
import org.truenewx.tnxjee.repo.mybatis.mapper.EntityMapper;

@Component
public class QueryMapperRegistry extends MybatisMapperRegistry {

    private final Configuration configuration;

    public QueryMapperRegistry(SqlSessionFactory sqlSessionFactory) {
        super(sqlSessionFactory.getConfiguration());
        this.configuration = sqlSessionFactory.getConfiguration();
        Map<Class<?>, MybatisMapperProxyFactory<?>> knownMappers = getKnownMappers(
                this.configuration.getMapperRegistry());
        if (knownMappers != null && !knownMappers.isEmpty()) {
            BeanUtil.setFieldValue(this, "knownMappers", knownMappers);
        }
        BeanUtil.setFieldValue(this.configuration, "mapperRegistry", this);
        BeanUtil.setFieldValue(this.configuration, "mybatisMapperRegistry", this);
    }

    private static Map<Class<?>, MybatisMapperProxyFactory<?>> getKnownMappers(MapperRegistry mapperRegistry) {
        return BeanUtil.getFieldValue(mapperRegistry, "knownMappers");
    }

    @Override
    public <T> void addMapper(Class<T> type) {
        if (type.isInterface()) {
            if (hasMapper(type)) {
                return;
            }
            boolean loadCompleted = false;
            Map<Class<?>, MybatisMapperProxyFactory<?>> knownMappers = getKnownMappers(this);
            try {
                MybatisMapperProxyFactory<T> mapperProxyFactory = new MybatisMapperProxyFactory<>(type);
                if (EntityMapper.class.isAssignableFrom(mapperProxyFactory.getMapperInterface())) {
                    Object methodCache = mapperProxyFactory.getMethodCache();
                    Object newCache = BeanUtil.createProxy(methodCache,
                            ((proxy, method, args) -> computeMethodCacheIfAbsent(methodCache, method, args)));
                    BeanUtil.setFieldValue(mapperProxyFactory, "methodCache", newCache);
                }
                knownMappers.put(type, mapperProxyFactory);

                Configuration config = BeanUtil.getFieldValue(this, "config");
                MybatisMapperAnnotationBuilder parser = new QueryMapperAnnotationBuilder(config, type);
                parser.parse();
                loadCompleted = true;
            } finally {
                if (!loadCompleted) {
                    knownMappers.remove(type);
                }
            }
        }
    }

    @SuppressWarnings("unchecked")
    private Object computeMethodCacheIfAbsent(Object methodCache, Method method, Object[] args)
            throws InvocationTargetException, IllegalAccessException {
        if ("computeIfAbsent".equals(method.getName())) {
            Method mapperInterfaceMethod = (Method) args[0];
            if (!mapperInterfaceMethod.isDefault()) {
                Class<?>[] parameterTypes = mapperInterfaceMethod.getParameterTypes();
                if (parameterTypes.length == 1 && Querying.class.isAssignableFrom(parameterTypes[0])
                        && QueryResult.class.isAssignableFrom(mapperInterfaceMethod.getReturnType())) {
                    Function<Method, Object> func = (Function<Method, Object>) args[1];
                    args[1] = (Function<Method, Object>) m -> {
                        Class<?> mapperInterface = mapperInterfaceMethod.getDeclaringClass();
                        QueryMapperMethod mapperMethod = new QueryMapperMethod(mapperInterface, mapperInterfaceMethod,
                                this.configuration);
                        Object plainMethodInvoker = func.apply(m);
                        BeanUtil.setFieldValue(plainMethodInvoker, "mapperMethod", mapperMethod);
                        return plainMethodInvoker;
                    };
                }
            }
        }
        return method.invoke(methodCache, args);
    }

}
